# -*- coding: utf-8 -*-
# nlp.py
# Created by Hardy on 21st, Jan
# Copyright 2017 杭州网川教育科技有限公司. All rights reserved.

from word2vec.word2vec import Word2VecEval
from text_classifier.classifier import Classifier
import jieba
import jieba.analyse as jana


class NLPService:
    def __init__(self, model_file, user_dict, stopwords_file, keyword_method, dv_model_path, lg_model_path,
                 tfidf_model_path=None):
        self.word2vec = Word2VecEval(model_file)
        self.text_classifier = Classifier(dv_model_path, lg_model_path, tfidf_model_path)
        self.method = keyword_method
        jieba.load_userdict(user_dict)
        with open(stopwords_file) as fd:
            stopwords = [l.strip() for l in fd.readlines() if l.strip()]
        self.stopword_dict = dict(zip(stopwords, [1] * len(stopwords)))

    def get_keywords(self, sentence, top_k=10, allow_pos=()):
        if len(sentence) < 50:
            if len(sentence) == len(sentence.encode('utf-8')):
                return sentence.split(' ')
        if self.method == 'tfidf':
            keywords = jana.extract_tags(sentence, topK=top_k, allowPOS=allow_pos)
        else:
            keywords = jana.textrank(sentence, topK=top_k, allowPOS=allow_pos)
        if len(keywords) == 0:
            keywords = self.get_seg(sentence)
        return keywords

    def get_similar_words(self, words, topn):
        return {"word2vec": self.word2vec.eval(words, topn)}

    def get_word_vectors(self, texts):
        out = []
        for t in texts:
            words = jieba.lcut(t.strip().lower())
            out.append(self.word2vec.get_vectors(words))
        return {'word2vec': out}

    def classify(self, keyword_list):
        return {"classify": self.text_classifier.classify(keyword_list)}

    def get_seg(self, sentence):
        return [w for w in jieba.cut(sentence) if w not in self.stopword_dict]
